#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <dirent.h>
#include <sys/stat.h>
#include <unistd.h>

#define MAX_FILENAME_LENGTH 256
#define RLE_MAGIC_NUMBER 0x1A5A

typedef struct {
    char filename[MAX_FILENAME_LENGTH];
    long offset;  //偏移量
    long compressed_size;
    long original_size;
    int is_directory;     //一个标志位，用于判断该条目是文件还是目录 真 目录 假 文件
} FileHeader;

typedef struct {
    int magic_number;  //魔数 若魔数不匹配会终止解压缩程序
    int file_count;
    time_t timestamp;   //时间戳
} ArchiveHeader;

// 修复的RLE压缩实现
long rle_compress(FILE *in, FILE *out) {
    int current_char, next_char;
    int count = 0;
    long compressed_size = 0;

    // 读取第一个字符
    current_char = fgetc(in);
    if (current_char == EOF) return 0;
    
    count = 1;

    while (1) {
        next_char = fgetc(in);
        
        if (next_char == current_char && count < 255) {
            // 连续重复字符，增加计数
            count++;
        } else {
            // 处理当前字符序列
            if (count >= 3) {
                // 使用RLE编码 (0xFF + 计数 + 字符)
                fputc(0xFF, out);
                fputc(count, out);
                fputc(current_char, out);
                compressed_size += 3;
            } else {
                // 不满足RLE条件，直接写入原始数据
                for (int i = 0; i < count; i++) {
                    fputc(current_char, out);
                    compressed_size++;
                }
            }
            
            // 重置计数
            count = 1;
            
            // 处理下一个字符
            if (next_char == EOF) break;
            current_char = next_char;
        }
    }
    
    return compressed_size;
}

// 递归处理目录
void process_directory(const char *dir_path, const char *base_path, FILE *archive, FileHeader **headers, int *file_count, int *current_capacity, long *total_original, long *total_compressed) {
    DIR *dir = opendir(dir_path);
    if (!dir) return;

    struct dirent *entry;
    struct stat statbuf;
    char full_path[MAX_FILENAME_LENGTH];
    char relative_path[MAX_FILENAME_LENGTH];

    while ((entry = readdir(dir)) != NULL) {
        if (strcmp(entry->d_name, ".") == 0 || strcmp(entry->d_name, "..") == 0)
            continue;

        snprintf(full_path, sizeof(full_path), "%s/%s", dir_path, entry->d_name);
        if (stat(full_path, &statbuf) == -1) continue;

        // 计算相对路径
        if (strlen(base_path) > 0 && strncmp(full_path, base_path, strlen(base_path)) == 0) {
            strcpy(relative_path, full_path + strlen(base_path));
            if (relative_path[0] == '/') memmove(relative_path, relative_path + 1, strlen(relative_path));
        } else {
            strncpy(relative_path, full_path, MAX_FILENAME_LENGTH - 1);
        }

        // 检查是否需要扩展内存
        if (*file_count >= *current_capacity) {
            *current_capacity *= 2;
            *headers = realloc(*headers, *current_capacity * sizeof(FileHeader));
        }

        (*headers)[*file_count].is_directory = S_ISDIR(statbuf.st_mode);
        strncpy((*headers)[*file_count].filename, relative_path, MAX_FILENAME_LENGTH - 1);

        if (S_ISDIR(statbuf.st_mode)) {
            (*headers)[*file_count].offset = ftell(archive);
            (*headers)[*file_count].original_size = 0;
            (*headers)[*file_count].compressed_size = 0;
            (*file_count)++;
            process_directory(full_path, base_path, archive, headers, file_count, current_capacity, total_original, total_compressed);
        } else {
            FILE *file = fopen(full_path, "rb");
            if (!file) continue;

            (*headers)[*file_count].offset = ftell(archive);
            (*headers)[*file_count].original_size = statbuf.st_size;

            long current_pos = ftell(archive);
            fseek(archive, sizeof(FileHeader), SEEK_CUR);

            long compressed_size = rle_compress(file, archive);
            fclose(file);

            (*headers)[*file_count].compressed_size = compressed_size;
            (*file_count)++;

            // 更新统计信息
            *total_original += (*headers)[*file_count-1].original_size;
            *total_compressed += compressed_size;

            // 写实际的文件头
            long temp_pos = ftell(archive);
            fseek(archive, current_pos, SEEK_SET);
            fwrite(&(*headers)[*file_count - 1], sizeof(FileHeader), 1, archive);
            fseek(archive, temp_pos, SEEK_SET);
        }
    }
    closedir(dir);
}
//写控制台
int main(int argc, char *argv[]) {
    if (argc < 3) {
        printf("用法: %s <输入文件/目录1> [输入文件/目录2 ...] <输出文件.rzp>\n", argv[0]);
        return 1;
    }

    clock_t start_time = clock();

    char *output_file = argv[argc - 1];
    int file_count = 0;
    int current_capacity = 100;
    long total_original_size = 0;
    long total_compressed_size = 0;

    FileHeader *headers = malloc(current_capacity * sizeof(FileHeader));
    FILE *archive = fopen(output_file, "wb");

    // 写入临时文件头
    ArchiveHeader temp_header = {RLE_MAGIC_NUMBER, 0, time(NULL)};
    fwrite(&temp_header, sizeof(ArchiveHeader), 1, archive);

    // 处理每个输入文件/目录
    for (int i = 1; i < argc - 1; i++) {
        struct stat statbuf;
        if (stat(argv[i], &statbuf) == -1) continue;

        if (S_ISDIR(statbuf.st_mode)) {
            if (file_count >= current_capacity) {
                current_capacity *= 2;
                headers = realloc(headers, current_capacity * sizeof(FileHeader));
            }

            headers[file_count].is_directory = 1;
            char relative_dir[MAX_FILENAME_LENGTH];
            char *last_slash = strrchr(argv[i], '/');
            strcpy(relative_dir, last_slash ? last_slash + 1 : argv[i]);
            
            strncpy(headers[file_count].filename, relative_dir, MAX_FILENAME_LENGTH - 1);
            headers[file_count].offset = ftell(archive);
            headers[file_count].original_size = 0;
            headers[file_count].compressed_size = 0;
            file_count++;

            process_directory(argv[i], argv[i], archive, &headers, &file_count, &current_capacity, &total_original_size, &total_compressed_size);
        } else {
            if (file_count >= current_capacity) {
                current_capacity *= 2;
                headers = realloc(headers, current_capacity * sizeof(FileHeader));
            }

            FILE *file = fopen(argv[i], "rb");
            if (!file) continue;

            headers[file_count].is_directory = 0;
            char *filename = strrchr(argv[i], '/');
            strncpy(headers[file_count].filename, filename ? filename + 1 : argv[i], MAX_FILENAME_LENGTH - 1);

            headers[file_count].offset = ftell(archive);
            headers[file_count].original_size = statbuf.st_size;

            long current_pos = ftell(archive);
            fseek(archive, sizeof(FileHeader), SEEK_CUR);

            long compressed_size = rle_compress(file, archive);
            fclose(file);

            headers[file_count].compressed_size = compressed_size;
            file_count++;

            // 更新统计信息
            total_original_size += headers[file_count-1].original_size;
            total_compressed_size += compressed_size;

            // 写入实际的文件头
            long temp_pos = ftell(archive);
            fseek(archive, current_pos, SEEK_SET);
            fwrite(&headers[file_count - 1], sizeof(FileHeader), 1, archive);
            fseek(archive, temp_pos, SEEK_SET);
        }
    }

    // 更新文件头
    fseek(archive, 0, SEEK_SET);
    temp_header.file_count = file_count;
    fwrite(&temp_header, sizeof(ArchiveHeader), 1, archive);

    // 写入所有文件头
    fseek(archive, sizeof(ArchiveHeader), SEEK_SET);
    fwrite(headers, sizeof(FileHeader), file_count, archive);

    fclose(archive);
    free(headers);

    clock_t end_time = clock();
    double elapsed_time = (double)(end_time - start_time) / CLOCKS_PER_SEC;

    printf("压缩完成！\n");
    printf("处理文件数: %d\n", file_count);
    printf("总原始大小: %.2f MB\n", total_original_size / (1024.0 * 1024.0));
    printf("总压缩大小: %.2f MB\n", total_compressed_size / (1024.0 * 1024.0));
    printf("压缩比: %.2f%%\n", (double)total_compressed_size / total_original_size * 100);
    printf("耗时: %.2f 秒\n", elapsed_time);

    return 0;
}    
